# encoding=utf-8
import os
import json
import shutil
import pandas as pd

BASE_PATH = r'D:\data'  # 原始文件夹路径
OUT_PATH = r'D:\data_result'  # 输出文件夹路径

BASE_NAME = BASE_PATH.rsplit('\\', 1)[-1]

JSON_FILE_NAME = 'dataset.json'
CSV_FILE_NAME = 'dataset.csv'

BASE_SUB_FOLDER_NAME1 = 'image'
BASE_SUB_FOLDER_NAME2 = 'label'
BASE_SUB_FOLDER_NAME3 = 'raw'

OUT_MAIN_FOLDER_NAME = 'save'
OUT_SUB_FOLDER_NAME1 = 'imagesTr'
OUT_SUB_FOLDER_NAME2 = 'labelsTr'
OUT_SUB_FOLDER_NAME3 = 'imagesTs'
OUT_SUB_FOLDERS = [OUT_SUB_FOLDER_NAME1, OUT_SUB_FOLDER_NAME2, OUT_SUB_FOLDER_NAME3]

NUMBER = 5
NUMBER_raw = 5
TEST_FOLDER = r'D:\data\raw'
TEST_PATH = os.path.dirname(__file__)
folder_path = os.path.join(TEST_PATH, TEST_FOLDER)
folder_path_list = os.listdir(folder_path)


def start():
    make_output_dir()
    handle_base_file()


def make_output_dir():
    if not os.path.exists(OUT_PATH):
        print(f'路径{OUT_PATH}不存在')
        return
    for name in OUT_SUB_FOLDERS:
        out_sub_folder = os.path.join(OUT_PATH, name)
        if os.path.exists(out_sub_folder):
            print(f'输出文件夹{out_sub_folder}已存在')
            continue
        os.mkdir(out_sub_folder)


def handle_base_file():
    if not os.path.exists(BASE_PATH):
        print(f'路径{BASE_PATH}不存在')
        return
    base_path = os.path.abspath(BASE_PATH)

    image_files = os.listdir(os.path.join(base_path, BASE_SUB_FOLDER_NAME1))
    label_files = os.listdir(os.path.join(base_path, BASE_SUB_FOLDER_NAME2))
    raw_files = os.listdir(os.path.join(base_path, BASE_SUB_FOLDER_NAME3))

    copy_base_file(image_files, label_files, raw_files)


def copy_base_file(image_files, label_files, raw_files):
    total = total_raw_counts = total_training_counts = 0
    base_path = os.path.abspath(BASE_PATH)
    out_path = os.path.abspath(OUT_PATH)
    csv_data = {'input': [], 'output': []}
    for img_name, label_name in zip(image_files, label_files):
        count = f'%0{NUMBER}d'
        count = count % total

        base_img_path = os.path.join(base_path, BASE_SUB_FOLDER_NAME1, img_name)
        out_img_name = f'{BASE_NAME}_{count}_0000.nii.gz'
        out_img_path = os.path.join(out_path, OUT_SUB_FOLDER_NAME1, out_img_name)
        csv_data['input'].append(img_name)
        csv_data['output'].append(out_img_name)

        base_label_path = os.path.join(base_path, BASE_SUB_FOLDER_NAME2, label_name)
        out_label_name = f'{BASE_NAME}_{count}.nii.gz'
        out_label_path = os.path.join(out_path, OUT_SUB_FOLDER_NAME2, out_label_name)
        csv_data['input'].append(label_name)
        csv_data['output'].append(out_label_name)

        shutil.copy(base_img_path, out_img_path)
        shutil.copy(base_label_path, out_label_path)
        print(f'{img_name}、{label_name}已处理')
        total_training_counts += 1
        total += 1


    for raw_name in raw_files:
        count = f'%0{NUMBER_raw}d'
        count = count % total
        base_raw_path = os.path.join(base_path, BASE_SUB_FOLDER_NAME3, raw_name)
        out_raw_name = f'{BASE_NAME}_{count}_0000.nii.gz'
        out_raw_path = os.path.join(out_path, OUT_SUB_FOLDER_NAME3, out_raw_name)
        csv_data['input'].append(raw_name)
        csv_data['output'].append(out_raw_name)
        shutil.copy(base_raw_path, out_raw_path)
        print(f'{raw_name}已处理')
        total_raw_counts += 1
        total += 1


    save_to_csv(csv_data)
    save_to_json(total_training_counts, total_raw_counts)


def save_to_json(total_training_counts, total_raw_counts):
    json_data = {
        'name': BASE_NAME,
        'description': 'Green',
        'tensorImageSize': '3D',
        'reference': 'Wall Motion data for nnunet',
        'licence': '',
        'release': '0.0',
        'modality': {
            '0': 'MR',
        },
        'labels': {
            '0': 'background',
            '1': 'Green',
        },
        'numTraining': total_training_counts,
        'training': [
            {
                'image': f'./{OUT_SUB_FOLDER_NAME1}/{BASE_NAME}_{name:0{NUMBER}}.nii.gz',
                'label': f'./{OUT_SUB_FOLDER_NAME2}/{BASE_NAME}_{name:0{NUMBER}}.nii.gz',
            }
            for name in range(total_training_counts)
        ],
        'numTest': total_raw_counts,
        'test': [
            f'./{OUT_SUB_FOLDER_NAME3}/{BASE_NAME}_{name:0{NUMBER_raw}}_0000.nii.gz'
            for name in range(total_training_counts, total_training_counts + total_raw_counts)
        ],
    }

    json_path = os.path.join(os.path.abspath(OUT_PATH), JSON_FILE_NAME)
    with open(json_path, 'w', encoding='utf-8') as f:
        f.write(json.dumps(json_data, indent=4))
    json_file = json_path.replace("\\", "/")
    print(f'json文件已保存到"file:////{json_file}"')


def save_to_csv(data):
    df = pd.DataFrame(data)
    csv_path = os.path.join(os.path.abspath(OUT_PATH), CSV_FILE_NAME)
    df.to_csv(csv_path, index=False)
    print(f'csv文件已保存到:{csv_path}')


if __name__ == '__main__':
    start()
